/*
 * Environment:
 * Microsoft Visual Studio 2008 Version 9.0.21022.8 RTM
 * Microsoft (R) 32-bit C/C++ Optimizing Compiler Version 15.00.21022.08 for 80x86
 *
 * Authors: Sergey Zamyslov, Sergey Moshkov
 */

#include <stdio.h>
#include <time.h>

#include <oclUtils.h>

#define DEBUG true

#define ROWS_PER_BLOCK 32
#define LOCAL_WORK_SIZE 256

// OpenCL kernels source code.
const char* OPENCL_SOURCE_FILE = "TDMA_OpenCL.cl";

const char* OPENCL_SOURCE_KERNEL_FORWARD = "forward";
const char* OPENCL_SOURCE_KERNEL_BACKWARD = "backward";
const char* OPENCL_SOURCE_KERNEL_AUXILIARY_MATRIX = "auxiliaryMatrixCoeffs";
const char* OPENCL_SOURCE_KERNEL_SET_X = "setX";
const char* OPENCL_SOURCE_KERNEL_X = "calculateX";

float function(float x) {
    return cosf(x) / x;
}

void initializeKnotsAndNodes(float a, float b,
                             float* knots, int knotCount,
                             float* nodesX, float* nodesY)
{
    int    i;
    float step = (b - a) / (knotCount - 1);

    float d = step * 0.9f;
    float dOver2 = d / 2.0f;

    knots[0]             = a;
    knots[knotCount - 1] = b;

    nodesX[0] = knots[0];
    nodesY[0] = function(nodesX[0]);

    srand((int)time(NULL));

    for (i = 1; i < knotCount - 1; i++)
    {
        knots[i] = i * step;
        knots[i] += (((float)rand() / RAND_MAX) * d) - dOver2;

        nodesX[i] = (knots[i] + knots[i - 1]) / 2;
        nodesY[i] = function(nodesX[i]);
    }

    i = knotCount - 1;
    nodesX[i] = (knots[i] + knots[i - 1]) / 2;
    nodesY[i] = function(nodesX[i]);

    nodesX[knotCount] = knots[knotCount - 1];
    nodesY[knotCount] = function(nodesX[knotCount]);
}

void solveTDM(int equationCount,
              float* subDiagonal, float* mainDiagonal, float* supDiagonal,
              float* b,
              float* x)
{
    // Initialization
    int    i;
    float* mainDiagonalWork = new float[equationCount],
         * bWork            = new float[equationCount];

	mainDiagonalWork[0] = mainDiagonal[0];
	bWork[0]            = b[0];

    // Forward
    for (i = 1; i < equationCount; i++)
    {
        float m = subDiagonal[i-1] / mainDiagonalWork[i-1];

        mainDiagonalWork[i] = mainDiagonal[i] - m * supDiagonal[i-1];
        bWork[i]            = b[i] - m * bWork[i-1];
	}

    // Backward
    x[equationCount - 1] = bWork[equationCount - 1] / mainDiagonalWork[equationCount - 1];	
    for (i = equationCount - 2; i >= 0; i--)
    {
        x[i] = (bWork[i] - supDiagonal[i] * x[i + 1]) / mainDiagonalWork[i];
    }

    delete[] mainDiagonalWork;
    delete[] bWork;
}



#define checkError(clError) if (clError != CL_SUCCESS) { printf("Error: line %u.\n\n", __LINE__); break; }

double solveTDM_Block(int equationCount,
                      float* subDiagonal, float* mainDiagonal, float* supDiagonal,
                      float* b,
                      float* x)
{
    // Check input.
    if(equationCount % ROWS_PER_BLOCK != 0)
    {
        printf("equationCount cannot be divided evenly by %d (ROWS_PER_BLOCK).", ROWS_PER_BLOCK);
        return 0.0;
    }

    int blockCount = equationCount / ROWS_PER_BLOCK;

    if(blockCount % LOCAL_WORK_SIZE != 0)
    {
        printf("blockCount cannot be divided evenly by %d (LOCAL_WORK_SIZE).", LOCAL_WORK_SIZE);
        return 0.0;
    }

    size_t clGlobalWorkSize = blockCount;
    size_t clLocalWorkSize = LOCAL_WORK_SIZE;

    if(DEBUG)
    {
        printf("Equation count = %d.\n", equationCount);
        printf("Rows per block = %d.\n", ROWS_PER_BLOCK);
        printf("Block count = %d.\n", blockCount);
        printf("\n");

        printf("Global work size = %d.\n", clGlobalWorkSize);
        printf("Local work size = %d.\n", clLocalWorkSize);
        printf("Work group count = %d.\n", clGlobalWorkSize / clLocalWorkSize);
        printf("\n");
    }

    cl_int clError;
    cl_platform_id clPlatform;
    cl_device_id clDevice;
    cl_context clContext = NULL;
    cl_command_queue clCommandQueue = NULL;

    cl_mem clDeviceSubDiagonal = NULL;
    cl_mem clDeviceMainDiagonal = NULL;
    cl_mem clDeviceSupDiagonal = NULL;
    cl_mem clDeviceB = NULL;
    cl_mem clDeviceX = NULL;
    cl_mem clDeviceAuxiliaryMatrix = NULL;

    char* clSource = NULL;
    size_t clSourceLength;

    cl_program clProgram = NULL;

    cl_kernel clKernelForward = NULL;
    cl_kernel clKernelBackward = NULL;
    cl_kernel clKernelAuxiliaryMatrix = NULL;
    cl_kernel clKernelSetX = NULL;
    cl_kernel clKernelX = NULL;

    float* auxiliaryMatrixSet = new float[blockCount * 5];
    float* auxiliaryMatrix[5];
    auxiliaryMatrix[0] = auxiliaryMatrixSet + 0 * blockCount;
    auxiliaryMatrix[1] = auxiliaryMatrixSet + 1 * blockCount;
    auxiliaryMatrix[2] = auxiliaryMatrixSet + 2 * blockCount;
    auxiliaryMatrix[3] = auxiliaryMatrixSet + 3 * blockCount;;
    auxiliaryMatrix[4] = auxiliaryMatrixSet + 4* blockCount;;

    double time = -1.0;

    do
    {
        /* OpenCL initialization. */
        clError = clGetPlatformIDs(1, &clPlatform, NULL);
        checkError(clError);

        clError = clGetDeviceIDs(clPlatform, CL_DEVICE_TYPE_GPU, 1, &clDevice, NULL);
        checkError(clError);

        clContext = clCreateContext(0, 1, &clDevice, NULL, NULL, &clError);
        checkError(clError);

        clCommandQueue = clCreateCommandQueue(clContext, clDevice, 0, &clError);
        checkError(clError);

        /* Allocate memory on device. */
        size_t dataElementSize = sizeof(cl_float);

        clDeviceSubDiagonal = clCreateBuffer(clContext, CL_MEM_READ_WRITE, dataElementSize * (equationCount - 1), NULL, &clError);
        checkError(clError);

        clDeviceMainDiagonal = clCreateBuffer(clContext, CL_MEM_READ_WRITE, dataElementSize * equationCount, NULL, &clError);
        checkError(clError);

        clDeviceSupDiagonal = clCreateBuffer(clContext, CL_MEM_READ_WRITE, dataElementSize * (equationCount - 1), NULL, &clError);
        checkError(clError);

        clDeviceB = clCreateBuffer(clContext, CL_MEM_READ_WRITE, dataElementSize * equationCount, NULL, &clError);
        checkError(clError);

        clDeviceX = clCreateBuffer(clContext, CL_MEM_READ_WRITE, dataElementSize * equationCount, NULL, &clError);
        checkError(clError);

        clDeviceAuxiliaryMatrix = clCreateBuffer(clContext, CL_MEM_READ_WRITE, dataElementSize * blockCount * 4, NULL, &clError);
        checkError(clError);

        /* OpenCL sources. */
        char clSourcePreambula[256];
        sprintf(clSourcePreambula,
            "#define ROWS_PER_BLOCK %d\n\n", ROWS_PER_BLOCK);

        clSource = oclLoadProgSource(OPENCL_SOURCE_FILE, clSourcePreambula, &clSourceLength);
        if(clSource == 0)
        {
            printf("Couldn't find OpenCL source file.\n\n");
            break;
        }

        if(DEBUG)
        {
            printf("OpenCL source:\n%s\n", clSource);
        }

        clProgram = clCreateProgramWithSource(clContext, 1, (const char**)&clSource, &clSourceLength, &clError);
        checkError(clError);

        clError = clBuildProgram(clProgram, 0, NULL, NULL, NULL, NULL);
        checkError(clError);

        /* Kernels. */
        clKernelForward = clCreateKernel(clProgram, OPENCL_SOURCE_KERNEL_FORWARD, &clError);
        checkError(clError);

        clError = clSetKernelArg(clKernelForward, 0, sizeof(cl_int), (void*)&equationCount);
        clError |= clSetKernelArg(clKernelForward, 1, sizeof(cl_mem), (void*)&clDeviceSubDiagonal);
        clError |= clSetKernelArg(clKernelForward, 2, sizeof(cl_mem), (void*)&clDeviceMainDiagonal);
        clError |= clSetKernelArg(clKernelForward, 3, sizeof(cl_mem), (void*)&clDeviceSupDiagonal);
        clError |= clSetKernelArg(clKernelForward, 4, sizeof(cl_mem), (void*)&clDeviceB);
        checkError(clError);

        clKernelBackward = clCreateKernel(clProgram, OPENCL_SOURCE_KERNEL_BACKWARD, &clError);
        checkError(clError);

        clError = clSetKernelArg(clKernelBackward, 0, sizeof(cl_int), (void*)&equationCount);
        clError |= clSetKernelArg(clKernelBackward, 1, sizeof(cl_mem), (void*)&clDeviceSubDiagonal);
        clError |= clSetKernelArg(clKernelBackward, 2, sizeof(cl_mem), (void*)&clDeviceMainDiagonal);
        clError |= clSetKernelArg(clKernelBackward, 3, sizeof(cl_mem), (void*)&clDeviceSupDiagonal);
        clError |= clSetKernelArg(clKernelBackward, 4, sizeof(cl_mem), (void*)&clDeviceB);
        checkError(clError);

        clKernelAuxiliaryMatrix = clCreateKernel(clProgram, OPENCL_SOURCE_KERNEL_AUXILIARY_MATRIX, &clError);
        checkError(clError);

        clError = clSetKernelArg(clKernelAuxiliaryMatrix, 0, sizeof(cl_int), (void*)&blockCount);
        clError |= clSetKernelArg(clKernelAuxiliaryMatrix, 1, sizeof(cl_mem), (void*)&clDeviceSubDiagonal);
        clError |= clSetKernelArg(clKernelAuxiliaryMatrix, 2, sizeof(cl_mem), (void*)&clDeviceMainDiagonal);
        clError |= clSetKernelArg(clKernelAuxiliaryMatrix, 3, sizeof(cl_mem), (void*)&clDeviceSupDiagonal);
        clError |= clSetKernelArg(clKernelAuxiliaryMatrix, 4, sizeof(cl_mem), (void*)&clDeviceB);
        clError |= clSetKernelArg(clKernelAuxiliaryMatrix, 5, sizeof(cl_mem), (void*)&clDeviceAuxiliaryMatrix);
        checkError(clError);

        clKernelSetX = clCreateKernel(clProgram, OPENCL_SOURCE_KERNEL_SET_X, &clError);
        checkError(clError);

        clError |= clSetKernelArg(clKernelSetX, 0, sizeof(cl_mem), (void*)&clDeviceAuxiliaryMatrix);
        clError |= clSetKernelArg(clKernelSetX, 1, sizeof(cl_mem), (void*)&clDeviceX);
        checkError(clError);

        clKernelX = clCreateKernel(clProgram, OPENCL_SOURCE_KERNEL_X, &clError);
        checkError(clError);

        clError = clSetKernelArg(clKernelX, 0, sizeof(cl_int), (void*)&equationCount);
        clError |= clSetKernelArg(clKernelX, 1, sizeof(cl_mem), (void*)&clDeviceSubDiagonal);
        clError |= clSetKernelArg(clKernelX, 2, sizeof(cl_mem), (void*)&clDeviceMainDiagonal);
        clError |= clSetKernelArg(clKernelX, 3, sizeof(cl_mem), (void*)&clDeviceSupDiagonal);
        clError |= clSetKernelArg(clKernelX, 4, sizeof(cl_mem), (void*)&clDeviceB);
        clError |= clSetKernelArg(clKernelX, 5, sizeof(cl_mem), (void*)&clDeviceX);
        checkError(clError);

        /* Load data. */
        clError = clEnqueueWriteBuffer(clCommandQueue, clDeviceSubDiagonal, CL_FALSE, 0, dataElementSize * (equationCount - 1), subDiagonal, 0, NULL, NULL);
        clError |= clEnqueueWriteBuffer(clCommandQueue, clDeviceMainDiagonal, CL_FALSE, 0, dataElementSize * equationCount, mainDiagonal, 0, NULL, NULL);
        clError |= clEnqueueWriteBuffer(clCommandQueue, clDeviceSupDiagonal, CL_FALSE, 0, dataElementSize * (equationCount - 1), supDiagonal, 0, NULL, NULL);
        clError |= clEnqueueWriteBuffer(clCommandQueue, clDeviceB, CL_FALSE, 0, dataElementSize * equationCount, b, 0, NULL, NULL);
        checkError(clError);



        /* Process data. */
        // Timing.
        clock_t start, finish;
        start = clock();

        clError = clEnqueueNDRangeKernel(clCommandQueue, clKernelForward, 1, NULL, &clGlobalWorkSize, &clLocalWorkSize, 0, NULL, NULL);
        checkError(clError);

        clError = clEnqueueNDRangeKernel(clCommandQueue, clKernelBackward, 1, NULL, &clGlobalWorkSize, &clLocalWorkSize, 0, NULL, NULL);
        checkError(clError);

        clError = clEnqueueNDRangeKernel(clCommandQueue, clKernelAuxiliaryMatrix, 1, NULL, &clGlobalWorkSize, &clLocalWorkSize, 0, NULL, NULL);
        checkError(clError);

        // Copy auxiliary matrix to host.
        clError = clEnqueueReadBuffer(clCommandQueue, clDeviceAuxiliaryMatrix, CL_TRUE, 0, dataElementSize * blockCount * 4, auxiliaryMatrixSet, 0, NULL, NULL);
        checkError(clError);

        solveTDM(blockCount,
            auxiliaryMatrix[0], auxiliaryMatrix[1], auxiliaryMatrix[2],
            auxiliaryMatrix[3],
            auxiliaryMatrix[4]);

        // Copy solution of auxiliary matrix back to device.
        clError = clEnqueueWriteBuffer(clCommandQueue, clDeviceAuxiliaryMatrix, CL_FALSE, 0, dataElementSize * blockCount, auxiliaryMatrix[4], 0, NULL, NULL);
        checkError(clError);

        clError = clEnqueueNDRangeKernel(clCommandQueue, clKernelSetX, 1, NULL, &clGlobalWorkSize, &clLocalWorkSize, 0, NULL, NULL);
        checkError(clError);

        clError = clEnqueueNDRangeKernel(clCommandQueue, clKernelX, 1, NULL, &clGlobalWorkSize, &clLocalWorkSize, 0, NULL, NULL);
        checkError(clError);



        // Copy data back to host.
        clError = clEnqueueReadBuffer(clCommandQueue, clDeviceX, CL_TRUE, 0, dataElementSize * equationCount, x, 0, NULL, NULL);
        checkError(clError);

        finish = clock();
        time = ((double)finish - (double)start) / CLOCKS_PER_SEC;
    }
    while(false);



    /* Clean up. */
    if(clSource)
    {
        delete[] clSource;
    }

    if(auxiliaryMatrixSet)
    {
        delete[] auxiliaryMatrixSet;
    }

    if(clKernelForward)
    {
        clReleaseKernel(clKernelForward);
    }

    if(clKernelBackward)
    {
        clReleaseKernel(clKernelBackward);
    }

    if(clKernelAuxiliaryMatrix)
    {
        clReleaseKernel(clKernelAuxiliaryMatrix);
    }

    if(clKernelSetX)
    {
        clReleaseKernel(clKernelSetX);
    }

    if(clKernelX)
    {
        clReleaseKernel(clKernelX);
    }

    if(clProgram)
    {
        clReleaseProgram(clProgram);
    }

    if(clCommandQueue)
    {
        clReleaseCommandQueue(clCommandQueue);
    }

    if(clContext)
    {
        clReleaseContext(clContext);
    }

    if(clDeviceSubDiagonal)
    {
        clReleaseMemObject(clDeviceSubDiagonal);
    }

    if(clDeviceMainDiagonal)
    {
        clReleaseMemObject(clDeviceMainDiagonal);
    }

    if(clDeviceSupDiagonal)
    {
        clReleaseMemObject(clDeviceSupDiagonal);
    }

    if(clDeviceB)
    {
        clReleaseMemObject(clDeviceB);
    }

    if(clDeviceX)
    {
        clReleaseMemObject(clDeviceX);
    }

    if(clDeviceAuxiliaryMatrix)
    {
        clReleaseMemObject(clDeviceAuxiliaryMatrix);
    }

    return time;
}

void calculateSplineCoeffs(float* knots, int knotCount, float* nodesY,
                           float* coeffs)
{
    // Prepare matrix.
    float* subDiagonal  = new float[knotCount - 1],
         * supDiagonal  = new float[knotCount - 1],
         * mainDiagonal = new float[knotCount],
         * b            = new float[knotCount];

    for (int i = 1; i < knotCount - 1; i++)
    {
        subDiagonal[i - 1] = knots[i] - knots[i - 1];
        supDiagonal[i - 1] = knots[i] - knots[i - 1];
        mainDiagonal[i]    = 3 * (knots[i + 1] - knots[i - 1]);
        b[i - 1]           = 8 * (nodesY[i] - nodesY[i - 1]);
    }

    mainDiagonal[0] = 3 * (knots[1] - knots[0]);
    mainDiagonal[knotCount - 1] = 3 * (knots[knotCount - 1] - knots[knotCount - 2]);
    subDiagonal[knotCount - 2]  = knots[knotCount - 1] - knots[knotCount - 2];
    supDiagonal[knotCount - 2]  = knots[knotCount - 1] - knots[knotCount - 2];
    b[knotCount - 2]            = 8 * (nodesY[knotCount - 1] - nodesY[knotCount - 2]);
    b[knotCount - 1]            = 8 * (nodesY[knotCount] - nodesY[knotCount - 1]);

    // Solve matrix.
    double time = solveTDM_Block(knotCount, subDiagonal, mainDiagonal, supDiagonal, b, coeffs);
    printf("Time: %f\n\n\n", time);

    delete[] subDiagonal;
    delete[] supDiagonal;
    delete[] mainDiagonal;
    delete[] b;
}

float getSplinePointY(int knotCount, float* knots,
                      float* nodesX, float* nodesY,
                      float* coeffs,
                      float x)
{
    int i;
    float c;

    for (i = 1; i < knotCount; ++i)
    {
        if (x <= knots[i])
        {
            break;
        }
    }

    if (i < knotCount)
    {
        c = x - nodesX[i];

        return nodesY[i] +
            0.5f * (coeffs[i] + coeffs[i - 1]) * c + 
            (coeffs[i] - coeffs[i - 1]) * c * c / (2 * (knots[i] - knots[i - 1]));
    }

    return 0.0f;
}

int main(int argc, char** argv)
{
    // Interpolation parameters.
    float intervalStart = 0.5f;
    float intervalEnd = 10.0f;
    int knotCount = 10 * 1024000;

    // Output parameters.
    bool output = true;
    bool outputNodes = false;
    char* outputFileName = "chart.csv";
    float chartDx = (intervalEnd - intervalStart) / 1000.0f;

    // Interpolation input data.
    float* knots  = new float[knotCount];
    float* nodesX = new float[knotCount + 1];
    float* nodesY = new float[knotCount + 1];

    // Interpolation output data.
    float* coeffs  = new float[knotCount];

    initializeKnotsAndNodes(intervalStart, intervalEnd, knots, knotCount, nodesX, nodesY);
    calculateSplineCoeffs(knots, knotCount, nodesY, coeffs);

    // Output.
    if(output)
    {
        FILE* outputFile = fopen(outputFileName, "w+");

        fprintf(outputFile, "X;Y;splineY;nodesY\n");
        float chartX;
        for(chartX = intervalStart; chartX < intervalEnd; chartX += chartDx) {
            fprintf(outputFile, "%f;%f;%f;NA\n",
                chartX,
                function(chartX),
                getSplinePointY(knotCount, knots, nodesX, nodesY, coeffs, chartX));
        }

        chartX = intervalEnd;
        fprintf(outputFile, "%f;%f;%f;NA\n", chartX,
            function(chartX),
            getSplinePointY(knotCount, knots, nodesX, nodesY, coeffs, chartX));

        if(outputNodes)
        {
            for(int i = 0; i <= knotCount; ++i) {
                fprintf(outputFile, "%f;NA;NA;%f\n",
                    nodesX[i],
                    nodesY[i]);
            }
        }

        fclose(outputFile);
    }

    delete[] coeffs;
    delete[] knots;
    delete[] nodesX;
    delete[] nodesY;

    return 0;
}
